# 本文件包含 1s 高斯轨道的相关积分
# 以及 STO3G 轨道的相关积分

from .molecules import *
import scipy.special as sp


def boys_function(x: float) -> float:
    '''
    0阶 boys function
    '''
    if x < 1e-6:
        return 1
    else:
        return 0.5 * np.sqrt(np.pi/x) * sp.erf(np.sqrt(x))


def overlap_integral(g1: Gaussion_s, g2: Gaussion_s) -> float:
    '''
    1s 高斯轨道之间的重叠积分
    '''
    p = g1.alpha + g2.alpha
    mu = g1.alpha * g2.alpha / p
    gamma = 4*mu / p
    return gamma**(3/4) * np.exp(-mu * (g1.center - g2.center)**2)


def kinetic_integral(g1: Gaussion_s, g2: Gaussion_s) -> float:
    '''
    1s 高斯轨道之间的动能积分
    '''
    p = g1.alpha + g2.alpha
    mu = g1.alpha * g2.alpha / p
    gamma = 4*mu / p
    R2 = (g1.center - g2.center) ** 2
    return gamma**(3/4) * mu * (3-2*mu*R2) * np.exp(-mu * R2)


def nuclear_attraction_integral(Zc: int, Rc: float, g1: Gaussion_s, g2: Gaussion_s) -> float:
    '''
    1s 高斯轨道之间的原子核吸引势积分
        Zc: 中心原子的电荷量
        Rc: 中心原子核的坐标
    '''
    p = g1.alpha + g2.alpha
    mu = g1.alpha * g2.alpha / p
    gamma = 4*mu / p
    Rmu = (g1.alpha * g1.center + g2.alpha * g2.center) / p
    R2 = (g1.center - g2.center) ** 2
    return -2 * Zc * gamma**(3/4) * np.sqrt(p/np.pi) * np.exp(-mu*R2) * \
        boys_function(p*(Rmu - Rc)**2)


def electron_repulsion_integral(g1: Gaussion_s, g2: Gaussion_s, g3: Gaussion_s, g4: Gaussion_s) -> float:
    '''
    1s 高斯轨道电子排斥积分
    '''
    p1 = g1.alpha + g4.alpha
    p2 = g2.alpha + g3.alpha
    mu1 = g1.alpha * g4.alpha / p1
    mu2 = g2.alpha * g3.alpha / p2
    gamma1 = 4 * mu1 / p1
    gamma2 = 4 * mu2 / p2
    Rmu1 = (g1.alpha * g1.center + g4.alpha * g4.center) / p1
    Rmu2 = (g2.alpha * g2.center + g3.alpha * g3.center) / p2
    lam = p1 * p2 / (p1 + p2)
    R = abs(Rmu1 - Rmu2)
    V = (gamma1*gamma2)**(3/4) * np.exp(-mu1 * (g1.center - g4.center)
                                        ** 2 - mu2 * (g2.center - g3.center)**2)
    if R < 1e-6:
        return V * 2 * np.sqrt(lam / np.pi)
    else:
        return V * sp.erf(np.sqrt(lam) * R) / R


def two_center_integral(s1: STO3G, s2: STO3G, integral) -> float:
    '''
    计算两个 STO3G 相关积分通用模板
    '''
    total = 0.0
    for i in range(3):
        for j in range(3):
            total += s1.d[i] * s2.d[j] * integral(s1.g[i], s2.g[j])
    return total


def four_center_integral(s1: STO3G, s2: STO3G, s3: STO3G, s4: STO3G, integral) -> float:
    '''
    计算4个 STO3G 相关积分通用模板
    '''
    total = 0.0
    for i in range(3):
        for j in range(3):
            for k in range(3):
                for l in range(3):
                    total += s1.d[i] * s2.d[j] * s3.d[k] * s4.d[l] * \
                        integral(s1.g[i], s2.g[j], s3.g[k], s4.g[l])
    return total


def overlap_integral_sto3g(s1: STO3G, s2: STO3G) -> float:
    '''
    两个 STO3G 轨道之间的重叠积分
    '''
    return two_center_integral(s1, s2, overlap_integral)


def kinetic_integral_sto3g(s1: STO3G, s2: STO3G) -> float:
    '''
    两个 STO3G 轨道之间的动能积分
    '''
    return two_center_integral(s1, s2, kinetic_integral)


def nuclear_attraction_integral_sto3g(Zc: int, Rc: float, s1: STO3G, s2: STO3G) -> float:
    '''
    两个 STO3G 轨道之间的原子核吸引势积分
        Zc: 中心原子的电荷量
        Rc: 中心原子核的坐标
    '''
    return two_center_integral(s1, s2,
                               lambda g1, g2: nuclear_attraction_integral(Zc, Rc, g1, g2))


def electron_repulsion_integral_sto3g(s1: STO3G, s2: STO3G, s3: STO3G, s4: STO3G) -> float:
    '''
    四个 STO3G 轨道之间的电子排斥积分
    '''
    return four_center_integral(s1, s2, s3, s4, electron_repulsion_integral)
